#pragma OPENCL EXTENSION cl_amd_printf : enable
#define MAX_NEIGH       96
#define N_UNROLL        4     //unroll, future by mako
#define SMALL           1.E-7f //avoid dividing by 0.

#define N_PER_CELL    ${natoms_per_cell} //bymako
#define NX  ${nx}  //by mako
#define NY  ${ny}  //by mako
#define NZ  ${nz}  //by mako
#define NCELLS  (NX*NY*NZ)

#define index3(i,j,k)    ((i*NY +j)*NZ + k)  //index of n_atoms_in_cell
#define index4(i,j,k,l)  (((i*NY + j)*NZ + k)*N_PER_CELL+l) //index of atoms_in_cell
#define swap(x,y,tmp)    ({tmp=x; x=y; y=tmp;}) 

//density cell: 
//n_atoms_in_cell[NX][NY][NZ]: n_atoms of a cell
//atoms_in_cell[NX][NY][NZ][N_PER_CELL]: atom-ids of a cell
//----------------------------------------------------
__kernel void init_cell(        //init to -1 and n_grid to 0
    __global int *atoms_in_cell,   //ids, NULL if -1
    __global int *n_atoms_in_cell  //
    ) 
{
    int iglb = get_global_id(0); 
    int sglb = get_global_size(0);

    int i = iglb;
    while(i<NCELLS){
        n_atoms_in_cell[i] = 0;
        for(int j=0;j<N_PER_CELL;++j){
            atoms_in_cell[i*N_PER_CELL+j] = -1;
        }
    //if(i<1) printf("ig:%d,init-\n",i);
        i += sglb;
    }

}
//---
__kernel void cal_cell( 
    const int n_atoms,
    const float4 box_lo,  //
    const float4 box_hi,  //
    __global float4 *r_t,  //[n_atoms] position, type
    __global float4 *rr_c, //[n_atoms]: s0-2:reduced unit, s3:cell-index
    __global int *atoms_in_cell,   //[NX][NY][NZ][N_PER_CELL] ids, NULL if -1
    __global int *n_atoms_in_cell, //[NX][NY][NZ] 
    __global int *overflow //output status flag, true(1) if overflow
    ) 
{
    int iglb = get_global_id(0); 
    int sglb = get_global_size(0);
    int iloc = get_local_id(0); 
    int nloc = get_local_size(0);
    float4 box_l = box_hi - box_lo;
    box_l.w = 1.f; //avoid dividing by 0
    int i = iglb;
    while(i<n_atoms){
        float4 ri = r_t[i]; //prefetch
        int igx,igy,igz,idx_ng,idx_g;
        float4 rr = fmod((ri-box_lo)/box_l,1.f); 
        rr = (rr<0)? 1.f-SMALL+rr : rr; //SMALL: avoid single precision overflow (eg. -0.0)
        igx = rr.x*NX;
        igy = rr.y*NY;
        igz = rr.z*NZ;
        idx_ng = index3(igx,igy,igz);
        rr.w = idx_ng;
        rr_c[i] = rr;
        __global volatile int *p = n_atoms_in_cell+idx_ng; 
        int ni = atomic_inc(p); //return previous value
        if(ni<N_PER_CELL){
            idx_g = index4(igx,igy,igz,ni);
            atoms_in_cell[idx_g] = i;
        }
        else{
            *overflow = 1;
            atomic_dec(p); //recover previous increase
        }
    //if(i<5) printf("atom-:%d,ni:%d\n",i,ni);
        ///
        i += sglb;
    }
}
//---
__kernel void sort_by_hilbert(
    const int n_atoms,
    __global float4 *in,  //[n_atoms] position, type
    __global float4 *out, //[n_atoms]: s0-2:reduced unit, s3:cell-index
    __global int *atoms_in_cell,   //[NX][NY][NZ][N_PER_CELL] ids, NULL if -1
    __global int *n_atoms_in_cell, //[NX][NY][NZ] 
    __global int *overflow //output status flag, true(1) if overflow
    ) 
{
}
//---
__kernel void create_neigh( 
    const int n_atoms,
    const float r2_cut,  //cutoff of r^2 
    const float4 box_l,  //boxsize, s3:should be no-zero to avoid dividing by 0
    const float4 box_p,  //box periodic, 1:periodic, 0:nonperiodic
    __global float4 *r_t,  //[n_atoms] position, type
    __global float4 *rr_c, //[n_atoms]: s0-2:reduced unit, s3:grid-index
    __global int *grid,   //[NX][NY][NZ][N_PER_CELL] ids, NULL if -1
    __global int *n_grid, //[NX][NY][NZ] num per grid
    __global int *neigh,   //[n_atoms][MAX_NEIGH] neighbor list, <0 means NULL
    __global int *n_neigh, //[n_atoms] real number of neighbor
    __global int *overflow //output status flag, true(1) if overflow, you may need larger MAX_NEIGH, please init it as false(0) before calling
    ) 
{
    const float sx = 1.f/NX; //reduced step in x
    const float sy = 1.f/NY; //reduced step in x
    const float sz = 1.f/NZ; //reduced step in x
    int iglb = get_global_id(0); 
    int sglb = get_global_size(0);
    int iloc = get_local_id(0); 
    int nloc = get_local_size(0);

    int i = iglb;
    while(i<n_atoms){
        int ni = 0; //n_neighbor of i
        float4 ri = r_t[i]; //prefetch
        float4 rr = rr_c[i];
        for(int j=0;j<MAX_NEIGH;++j){ //init all to -1
            int idx = i*MAX_NEIGH+j;
            neigh[idx] = -1;
        }

        
        //use mako to unroll:
        for(int ii=-1;ii<2;++ii){
            float fgx = rr.x + ii*sx;
            fgx = fmod(fgx,1.f);
            fgx = (fgx<0)? 1.f-SMALL+fgx : fgx; //SMALL: avoid single precision overflow (eg. -0.0)
            int jgx = fgx*NX;
            for(int jj=-1;jj<2;++jj){
                float fgy = rr.y + jj*sy;
                fgy = fmod(fgy,1.f);
                fgy = (fgy<0)? 1.f-SMALL+fgy : fgy; //SMALL: avoid single precision overflow (eg. -0.0)
                int jgy = fgy*NY;
                for(int kk=-1;kk<2;++kk){
                    float fgz = rr.z + kk*sz;
                    fgz = fmod(fgz,1.f);
                    fgz = (fgz<0)? 1.f-SMALL+fgz : fgz; //SMALL: avoid single precision overflow (eg. -0.0)
                    int jgz = fgz*NZ;
                    //
                    int idx_n_grid = index3(jgx,jgy,jgz);
                    for(int kj=0;kj<n_grid[idx_n_grid];++kj){
                        int idx_grid = index4(jgx,jgy,jgz,kj);
                        int j = grid[idx_grid];
                        float4 rj = r_t[j];
                        float4 d = ri - rj;
                        d = d - round(d/box_l)*box_l*box_p; //periodic or nonperiodic box
                        float r2 = d.x*d.x + d.y*d.y + d.z*d.z;
                        if (r2 < r2_cut && i != j){
                            int idx = i*MAX_NEIGH + ni;
                            neigh[idx] = j;
                            ++ni;
                            if(ni>=MAX_NEIGH){
                                *overflow = true;
                                //break;
                                goto over;
                            }
                        }
                    }
                }
            }
        }
over:
        n_neigh[i] = ni;
        i += sglb;
    }
}
//----------------------------------------------------
__kernel void create_neigh_slow( 
    const int n_atoms,
    const float r2_cut,  //cutoff of r^2 
    const float4 box_l,  //boxsize, s3:should be no-zero to avoid dividing by 0
    const float4 box_p,  //box periodic, 1:periodic, 0:nonperiodic
    __global float4 *r_t,  //[n_atoms] position, type
    __global int *neigh,   //[n_atoms][MAX_NEIGH] neighbor list, <0 means NULL
    __global int *n_neigh, //[n_atoms] real number of neighbor
    __global int *overflow //output status flag, true(1) if overflow, you may need larger MAX_NEIGH, please init it as false(0) before calling
    ) 
{
    int iglb = get_global_id(0); 
    int sglb = get_global_size(0);
    int iloc = get_local_id(0); 
    int nloc = get_local_size(0);

    int i = iglb;
    while(i<n_atoms){
        int ni = 0; //n_neighbor of i
        float4 xi = r_t[i]; //prefetch
        for(int j=0;j<MAX_NEIGH;++j){ //init all to -1
            int idx = i*MAX_NEIGH+j;
            neigh[idx] = -1;
        }
        int j = 0;
        while(j<n_atoms){
          //for(int kk=0;kk<N_UNROLL;kk++){ //let compiler do unrolling 
            float4 xj = r_t[j];
            float4 d = xj - xi;
            d = d - round(d/box_l)*box_l*box_p; //periodic or nonperiodic box
            float r2 = d.x*d.x + d.y*d.y + d.z*d.z;
            if (r2 < r2_cut && i != j){
                int idx = i*MAX_NEIGH + ni;
                neigh[idx] = j;
                ++ni;
                if(ni>=MAX_NEIGH){
                    *overflow = true;
                    //break;
                    goto over;
                }
            }
            ++j;
          //} //end-unroll
        } //end-while
        n_neigh[i] = ni;
        //if(i<4) printf("ig:%d,n-neigh:%d\n",i,ni);
        ///
over:
        i += sglb;
    }
}
